from typing import Union, Callable, List, Optional, Sequence

import paddle
from paddle.fluid.dygraph import Layer
from paddle.fluid.optimizer import Optimizer

from helm.dynamic.engine.engine import Engine
from helm.dynamic.engine.callback import DefaultTrainLogger, DefaultEvalLogger, ModelCheckpoint
from helm.dynamic.engine.metrics import Metric


def create_supervised_trainer(
        model: Layer,
        criterion: Union[Callable, Layer],
        optimizer: Optimizer,
        metrics: List[Metric],
        log_freq: int = 10,
        save_freq: Optional[int] = None,
        save_dir: Optional[int] = None,
        auxiliary_weight: Optional[float] = None,
):
    def train_batch(engine, batch):
        engine.model.train()

        image, label = batch

        logits = engine.model(image)
        if auxiliary_weight and isinstance(logits, Sequence) and len(logits) == 2:
            logits, logits_aux = logits
            loss = engine.criterion(logits, label) + auxiliary_weight * engine.criterion(logits_aux, label)
        else:
            loss = engine.criterion(logits, label)
        loss.backward()
        engine.optimizer.minimize(loss)
        engine.model.clear_gradients()

        output = {
            "loss": float(loss.numpy()),
            "y_true": label,
            "y_pred": logits,
            "batch_size": image.shape[0],
            "lr": engine.optimizer.current_step_lr(),
        }

        return output

    callbacks = [*metrics, DefaultTrainLogger(log_freq)]
    if save_freq:
        callbacks.append(ModelCheckpoint(save_freq, save_dir))
    engine = Engine(
        train_batch, callbacks,
        model=model, criterion=criterion, optimizer=optimizer)
    return engine


def create_supervised_evaluator(
        model: Layer,
        criterion: Union[Callable, Layer],
        metrics: List[Metric],
        log_freq: int = 10,
        stage: str = "Valid",
):

    def test_batch(engine, batch):
        engine.model.eval()

        image, label = batch

        with paddle.fluid.dygraph.no_grad():
            pred = engine.model(image)
            loss = engine.criterion(pred, label)

        output = {
            "loss": float(loss.numpy()),
            "y_true": label,
            "y_pred": pred,
            "batch_size": image.shape[0],
        }

        return output

    callbacks = [*metrics, DefaultEvalLogger(log_freq)]
    engine = Engine(test_batch, callbacks,
                    model=model, criterion=criterion, stage=stage)

    return engine


